"""Helpers for AWS."""
import argparse
import base64
import configparser
import os
import sys
import xml.etree.ElementTree as ET

import boto3
from cki_lib import yaml
from cki_lib.session import get_session
import requests_gssapi

SESSION = get_session('cki_tools.aws', raise_for_status=True)


def login(profile: str | None, account: str | None, role: str | None, duration: int | None) -> None:
    """Log into AWS via SAML."""
    auth = requests_gssapi.HTTPSPNEGOAuth(mutual_authentication=requests_gssapi.OPTIONAL)
    response = SESSION.get(os.environ['AWS_IDP_URL'], auth=auth)
    assertion = ET.fromstring(response.text).find(
        './/{*}input[@name="SAMLResponse"]'
    ).attrib['value']

    aws_roles = [e.text.split(',') for e in ET.fromstring(base64.b64decode(assertion)).findall(
        './/{*}Attribute[@Name="https://aws.amazon.com/SAML/Attributes/Role"]//{*}AttributeValue'
    )]
    role_arn, principal_arn = next((
        (r, p) for r, p in aws_roles if
        (not role or r.endswith(f'-{role}')) and
        (not account or f':{account}:' in p)
    ))

    token = boto3.client('sts').assume_role_with_saml(
        RoleArn=role_arn,
        PrincipalArn=principal_arn,
        SAMLAssertion=assertion,
        **({'DurationSeconds': duration} if duration else {}),
    )['Credentials']

    if profile:
        file_path = os.environ.get('AWS_SHARED_CREDENTIALS_FILE',
                                   os.path.expanduser('~/.aws/credentials'))
        conf = configparser.RawConfigParser()
        conf.read(file_path)

        if not conf.has_section(profile):
            conf.add_section(profile)
        conf.set(profile, 'aws_access_key_id', token['AccessKeyId'])
        conf.set(profile, 'aws_secret_access_key', token['SecretAccessKey'])
        conf.set(profile, 'aws_session_token', token['SessionToken'])

        file_dir = os.path.dirname(file_path)
        os.makedirs(file_dir, 0o700, exist_ok=True)
        with open(file_path, 'w+', encoding='utf8') as conf_file:
            conf.write(conf_file)
    else:
        print(yaml.dump(token))


def login_cli(args: list[str] | None = None) -> int:
    """Log into AWS via SAML."""
    return main(['login'] + (args if args is not None else sys.argv[1:]))


def main(args: list[str] | None = None) -> int:
    """Helpers for AWS."""
    parser = argparse.ArgumentParser(description='Helpers for AWS')

    subparsers = parser.add_subparsers(dest='type', required=True)

    parser_login = subparsers.add_parser(
        'login', help='Log into AWS via SAML', description='Log into AWS via SAML')
    parser_login.add_argument('--profile', help='AWS configuration profile')
    parser_login.add_argument('--account', help='AWS account')
    parser_login.add_argument('--role', help='role to assume')
    parser_login.add_argument('--duration', type=int, help='session duration')

    parsed_args = parser.parse_args(args)

    match parsed_args.type:
        case 'login':
            login(parsed_args.profile, parsed_args.account, parsed_args.role, parsed_args.duration)
            return 0


if __name__ == '__main__':
    sys.exit(main())
